#include "auxfns.h"
#include "complex_types.h"
#include "l1spirit_config.h"
#include <fftw3.h>
#include <omp.h>

void
rescale_data (cplxf *data, int *dsize, float& scale_factor,
              l1spirit_config const& config)
{
	int ds0 = dsize[0], ds1 = dsize[1], ds2 = dsize[2], nch = dsize[3];

	int *nzmask, nnz = sampling_pattern (nzmask, data, dsize);

	for (int i2 = 0; i2 < ds2; ++i2)
		for (int i1 = 0; i1 < ds1; ++i1)
		{
			int& nz = nzmask[i1 + ds1*i2];

			nz = 0;

			for (int i0 = ds0/2-10; i0 < ds0/2+10; ++i0)
				nz |= (abs(data[i0 + ds0*(i1 + ds1*i2)]) != 0);

			nnz += nz;
		}

	float avg_density = 1.0f*nnz / (ds1*ds2);
	int r = (int) ceil (1.0f / avg_density);
	if (r % 2) ++r;

	float *density = new float[ds1*ds2];

	for (int i2 = 0; i2 < ds2; ++i2)
		for (int i1 = 0; i1 < ds1; ++i1)
		{
			float& dens = density[i1 + ds1*i2];
			dens = 0;

			int nc = 0;
		

			for (int r1 = -r/2; r1 <= r/2; ++r1)
				for (int r2 = -r/2; r2 <= r/2; ++r2)
				{
					if (i1+r1 >= 0 && i1+r1 < ds1
					and i2+r2 >= 0 && i2+r2 < ds2){
						dens += nzmask [i1+r1 + ds1*(i2+r2)];
						++nc;
					}
				}

			if (0 == (dens *= 1.0f/nc))
				dens = avg_density;
		}


	l1spirit_debug(config,density,0,density,2,ds1,ds2);

	for (int i1 = 0; i1 < ds1*ds2*nch; ++i1)
		for (int i0 = 0; i0 < ds0; i0 += 2)
			data[i0 + ds0*i1] = -data[i0 + ds0*i1];

	int chsz; cplxf *ch = new cplxf[chsz = ds0*ds1*ds2];

	fftwf_iodim dims[3] = {
		{ds2, ds1*ds0, ds1*ds0},
		{ds1, ds0, ds0},
		{ds0, 1, 1}};

	fftwf_plan_with_nthreads (omp_get_max_threads());
	fftwf_plan p = fftwf_plan_guru_dft (3, dims, 0, 0,
	                                    (fftwf_complex*)ch, 
	                                    (fftwf_complex*)ch,
	                                     FFTW_BACKWARD, FFTW_ESTIMATE);

	scale_factor = -1.0f;

	for (int ich = 0; ich < nch; ++ich)
	{
		for (int i1 = 0; i1 < ds1*ds2; ++i1)
			for (int i0 = 0; i0 < ds0; ++i0)
				ch[i0 + ds0*i1] = data[i0 + ds0*i1 + ich*chsz]*(1./density[i1]);

		fftwf_execute (p);

		for (int i = 0; i < chsz; ++i)
			if (scale_factor < abs(ch[i]))
				scale_factor = abs(ch[i]);
	}


	fftwf_destroy_plan (p);
	scale_factor = scale_factor / sqrt(ds0*ds1*ds2);

	for (int i = 0; i < ds0*ds1*ds2*nch; ++i)
		data[i] *= (1.0 / scale_factor);
}

void
unscale_data (cplxf *data, int *dsize, float scale_factor)
{
	int ds0 = dsize[0], ds1 = dsize[1], ds2 = dsize[2], nch = dsize[3];

	for (int i = 0; i < ds0*ds1*ds2*nch; ++i)
		data[i] *= scale_factor;
}
